from celery import shared_task
from pymongo import MongoClient
from django.conf import settings

from es.models import ProblemIndex

# 连接 MongoDB
mongo_client = MongoClient(settings.MONGO_URI)
db = mongo_client[settings.MONGO_DB_NAME]
collection = db["tpa_online"]


@shared_task
def sync_mongo_to_es():
    """ 同步 MongoDB 数据到 Elasticsearch """
    for doc in collection.find({}, {"_id": 1, "problem": 1}):
        problem_doc = ProblemIndex(meta={"id": str(doc["_id"])}, id=str(doc["_id"]), problem=doc["problem"])
        problem_doc.save()
